/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: dispatcher aicpu
 */

#ifndef HCCL_DISPATCHER_AICPU_PUB_H
#define HCCL_DISPATCHER_AICPU_PUB_H

#include <vector>
#include <functional>
#include "sal_pub.h"
#include "dispatcher_pub.h"

#include "rtsq_interact/aicpu_hccl_sqcq.h"
#include "rtsq_interact/aicpu_hccl_sqcqv1.h"
#include "rtsq_interact/aicpu_hccl_sqcqv2.h"

namespace hccl {
using AddOneNotifyWaitSqe = void(*)(uint16_t, uint16_t, u64, const uint8_t *, uint8_t *, const dfx::DfxTimeOutConfig &);
using AddOneRecordSqe = void(*)(uint16_t, uint16_t, u64, const uint8_t *, uint8_t *);
using AddOneWriteValueRecordSqe = void(*)(uint16_t, uint16_t, u64, const uint8_t *, uint8_t *);
using AddOneMemcpySqe = void(*)(uint16_t, uint16_t, const void *, uint32_t, const rtDataType_t,
    rtRecudeKind_t, const void *, uint32_t, uint32_t, uint32_t, u64, uint8_t, const uint8_t *, uint8_t *);
using AddOneEventResetSqe = void(*)(uint16_t, int32_t, uint16_t, int64_t, int64_t,
    u64, const uint8_t *, uint8_t *);
using AddOneEventRecordSqe = void(*)(uint16_t, int32_t, uint16_t, const uint8_t *, uint8_t *);
using AddOneEventWaitSqe = void(*)(uint16_t, int32_t, uint16_t, const uint8_t *, uint8_t *);
using AddOneRdmaDbSendSqe = void(*)(uint16_t, uint16_t, uint64_t, uint64_t, uint32_t, uint8_t, const uint8_t *, uint8_t *);
using AddOnePlaceHolderSqe = void(*)(uint16_t, uint16_t, uint16_t, const uint8_t *, uint8_t *);

class DispatcherAiCpu : public DispatcherPub {
public:
    explicit DispatcherAiCpu(const u32 devPhyId);
    ~DispatcherAiCpu() override;
    HcclResult Init() override;
    u32 GetMaxNotifyWaitTime();
    HcclResult WaitValue(hccl::Stream &stream, u64 waitAddr, u64 valueAddr, bool reset) override;
    HcclResult WriteValue(hccl::Stream &stream, u64 writeAddr, u64 valueAddr) override;
    HcclResult SignalRecord(HcclRtNotify signal, hccl::Stream &stream, u32 userRank, u64 offset = INVALID_U64,
        s32 stage = INVALID_VALUE_STAGE, bool inchip = false, u64 signalAddr = INVALID_U64,
        u32 notifyId = INVALID_UINT) override;
    HcclResult SignalWait(HcclRtNotify signal, hccl::Stream &stream, u32 userRank, u32 remoteUserRank,
        s32 stage = INVALID_VALUE_STAGE, bool inchip = false, u32 notifyId = INVALID_UINT,
        u32 timeOut = NOTIFY_INVALID_WAIT_TIME) override;
    HcclResult MemcpyAsync(hccl::DeviceMem &dst, const hccl::DeviceMem &src, hccl::Stream &stream,
        u32 remoteUserRank = INVALID_VALUE_RANKID, hccl::LinkType inLinkType = hccl::LinkType::LINK_ONCHIP) override;
    HcclResult ReduceAsync(const void *src, void *dst, u64 dataCount, const HcclDataType datatype, HcclReduceOp redOp,
        Stream &stream, HcclReduceType reduceType = HcclReduceType::HCCL_TBE_REDUCE) override;
    HcclResult InlineReduceAsync(const void *src, u64 dataCount, const HcclDataType datatype, HcclReduceOp redOp,
        Stream &stream, void *dst, u32 remoteUserRank = INVALID_VALUE_RANKID,
        hccl::LinkType inLinkType = hccl::LinkType::LINK_ONCHIP) override;
    HcclResult LaunchTasksEx(Stream &stream, std::vector<Stream> &subStreams) override;
    HcclResult LaunchAllTasks() override;

    HcclResult RdmaSend(u32 dbindex, u64 dbinfo, const struct send_wr &wr, hccl::Stream &stream,
        u32 remoteUserRank = INVALID_VALUE_RANKID) override;
    HcclResult RdmaSend(u32 dbindex, u64 dbinfo, const struct send_wr &wr, hccl::Stream &stream, u32 userRank,
        u64 offset) override;

    HcclResult LaunchTask(Stream &stream, bool isBlockLaunch);
    HcclResult TbeReduceAsync(const void *src1, const void *src2, u64 count, const HcclDataType datatype,
        HcclReduceOp redOp, Stream &stream, const void *dst);
    HcclResult AddRetryPreamble(Stream &stream) override;
    void SetDbAddr(u64 dbAddr) override
    {
        dbAddr_ = dbAddr;
    }

    AddOneNotifyWaitSqe GetAddOneNotifyWaitSqe()
    {
        return addOneNotifyWaitSqe_;
    }
    AddOneRecordSqe GetAddOneRecordSqe()
    {
        return addOneRecordSqe_;
    }
    AddOneWriteValueRecordSqe GetAddOneWriteValueRecordSqe()
    {
        return addOneWriteValueRecordSqe_;
    }
    AddOneMemcpySqe GetAddOneMemcpySqe()
    {
        return addOneMemcpySqe_;
    }
    AddOneEventResetSqe GetAddOneEventResetSqe()
    {
        return addOneEventResetSqe_;
    }
    AddOneEventRecordSqe GetAddOneEventRecordSqe()
    {
        return addOneEventRecordSqe_;
    }
    AddOneEventWaitSqe GetAddOneEventWaitSqe()
    {
        return addOneEventWaitSqe_;
    }
    AddOneRdmaDbSendSqe GetAddOneRdmaDbSendSqe()
    {
        return addOneRdmaDbSendSqe_;
    }
    AddOnePlaceHolderSqe GetAddOnePlaceHolderSqe()
    {
        return addOneFlipPlaceHolderSqe_;
    }
    void SetOpExecStatusCallback(std::function<HcclResult()> checkOpExecStatusCallback)
    {
        checkOpExecStatusCallback_ = checkOpExecStatusCallback;
    }
public:
    dfx::DfxTimeOutConfig dfxTimeOutConfig_ = {0};
    uint32_t opRingBufferIdx_ = 0;
private:
    HcclResult AddFlipTask(Stream &stream);
    HcclResult RdmaSend(u32 dbIndex, u64 dbInfo, const struct send_wr &wr, hccl::Stream &stream,
        RdmaType rdmaType, u32 userRank, u64 offset = 0);
    HcclResult GetStreamSqeBufferAddr(hccl::Stream &stream, uint8_t *&sqeBufferAddr, uint8_t *&sqeTypeAddr,
        uint8_t *&sqeDfxInfoAddr, uint16_t &taskId);
    void SaveStreamInfo(hccl::Stream &stream);

    AddOneNotifyWaitSqe addOneNotifyWaitSqe_ = nullptr;
    AddOneRecordSqe addOneRecordSqe_ = nullptr;
    AddOneWriteValueRecordSqe addOneWriteValueRecordSqe_ = nullptr;
    AddOneMemcpySqe addOneMemcpySqe_ = nullptr;
    AddOneEventResetSqe addOneEventResetSqe_ = nullptr;
    AddOneEventRecordSqe addOneEventRecordSqe_ = nullptr;
    AddOneEventWaitSqe addOneEventWaitSqe_ = nullptr;
    AddOneRdmaDbSendSqe addOneRdmaDbSendSqe_ = nullptr;
    AddOnePlaceHolderSqe addOneFlipPlaceHolderSqe_ = nullptr;
    std::function<HcclResult()> checkOpExecStatusCallback_ = nullptr;

    u64 dbAddr_;
    HcclAicpuDispatcherInfo aicpuInfo_;

    std::unordered_map<s32, Stream> streamMap_; // 保存下过task的stream
};
} // namespace hccl
#endif // HCCL_DISPATCHER_AICPU_PUB_H